"""
Ymodem协议实现模块
用于MCU固件升级的Ymodem协议通信
"""

import time
import os
from PySide6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QGridLayout, 
                               QLabel, QLineEdit, QComboBox, QGroupBox, QCheckBox,
                               QSpinBox)
from PySide6.QtCore import QThread, Signal, QTimer
from PySide6.QtGui import QFont


class YmodemProtocol(QThread):
    """Ymodem协议处理线程"""
    
    # 信号定义
    status_update = Signal(str)          # 状态更新信号
    progress_update = Signal(int, int)   # 进度更新信号 (当前, 总数)
    data_send = Signal(bytes)            # 数据发送信号
    transfer_complete = Signal(bool, str)  # 传输完成信号 (成功/失败, 消息)
    ack_received = Signal(int)           # ACK接收信号 (响应类型)
    
    # Ymodem控制字符
    SOH = 0x01  # 128字节数据包头
    STX = 0x02  # 1024字节数据包头
    EOT = 0x04  # 传输结束
    ACK = 0x06  # 确认
    NAK = 0x15  # 否认
    CAN = 0x18  # 取消
    C   = 0x43  # 'C' - CRC请求
    
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.file_path = ""
        self.file_data = b""
        self.file_size = 0
        self.packet_size = 128  # 默认128字节包
        self.packet_number = 0
        self.total_packets = 0
        self.retry_count = 0
        self.max_retries = 3
        self.timeout = 1000  # 1秒超时
        self.running = False
        self.received_data = b""
        self.waiting_for_ack = False
        self.ack_response = None
        
        # 连接ACK接收信号到处理函数
        self.ack_received.connect(self.on_ack_received)
        
        # 应用传入的配置
        if config:
            self.set_config(config)
        
    def set_file_path(self, file_path):
        """设置要传输的文件路径"""
        self.file_path = file_path
    
    def set_file_data(self, file_data, file_name="firmware.bin"):
        """直接设置要传输的文件数据"""
        self.file_data = file_data
        self.file_size = len(file_data)
        self.file_path = file_name  # 用于显示
        self.total_packets = (self.file_size + self.packet_size - 1) // self.packet_size
        
    def set_config(self, config):
        """更新协议配置"""
        self.config = config
        self.packet_size = int(config.get('packet_size', 128))
        self.max_retries = int(config.get('max_retries', 3))
        self.timeout = int(config.get('timeout', 1000))
        
        # 调试输出配置信息
        print(f"DEBUG: 协议配置已更新 - 数据包大小: {self.packet_size}, 超时: {self.timeout}, 重试: {self.max_retries}")
        
    def load_file(self):
        """加载文件数据"""
        try:
            # 如果已经有文件数据，则不需要重新加载
            if self.file_data and self.file_size > 0:
                self.total_packets = (self.file_size + self.packet_size - 1) // self.packet_size
                
                # 调试输出文件分包信息
                print(f"DEBUG: 使用预设文件数据 - 文件大小: {self.file_size}字节, 包大小: {self.packet_size}字节, 总包数: {self.total_packets}")
                
                self.status_update.emit(f"文件数据准备就绪: {os.path.basename(self.file_path)} ({self.file_size} 字节)")
                return True
            
            # 从文件路径加载数据
            if not os.path.exists(self.file_path):
                raise FileNotFoundError(f"文件不存在: {self.file_path}")
                
            with open(self.file_path, 'rb') as f:
                self.file_data = f.read()
                
            self.file_size = len(self.file_data)
            self.total_packets = (self.file_size + self.packet_size - 1) // self.packet_size
            
            # 调试输出文件分包信息
            print(f"DEBUG: 文件分包 - 文件大小: {self.file_size}字节, 包大小: {self.packet_size}字节, 总包数: {self.total_packets}")
            
            self.status_update.emit(f"文件加载成功: {os.path.basename(self.file_path)} ({self.file_size} 字节)")
            return True
            
        except Exception as e:
            self.status_update.emit(f"文件加载失败: {str(e)}")
            return False
            
    def run(self):
        """协议主循环"""
        self.running = True
        self.status_update.emit(f"开始Ymodem协议传输... (数据包大小: {self.packet_size}字节)")
        
        # 加载文件
        if not self.load_file():
            self.transfer_complete.emit(False, "文件加载失败")
            return
            
        try:
            # 步骤1: 等待接收端准备就绪
            if not self.wait_for_receiver_ready():
                self.transfer_complete.emit(False, "接收端未准备就绪")
                return
                
            # 步骤2: 发送文件信息包
            if not self.send_file_info():
                self.transfer_complete.emit(False, "文件信息发送失败")
                return
                
            # 步骤3: 发送文件数据
            if not self.send_file_data():
                self.transfer_complete.emit(False, "文件数据传输失败")
                return
                
            # 步骤4: 发送传输结束信号
            if not self.send_end_of_transmission():
                self.transfer_complete.emit(False, "传输结束信号发送失败")
                return
                
            self.status_update.emit("Ymodem传输完成!")
            self.transfer_complete.emit(True, "传输成功完成")
            
        except Exception as e:
            self.status_update.emit(f"传输过程出错: {str(e)}")
            self.transfer_complete.emit(False, f"传输异常: {str(e)}")
            
    def wait_for_receiver_ready(self):
        """等待接收端准备就绪信号"""
        self.status_update.emit("等待接收端准备就绪...")
        self.retry_count = 0
        
        while self.retry_count < self.max_retries and self.running:
            # 设置等待准备就绪状态
            self.waiting_for_ack = True
            self.ack_response = None
            
            # 等待接收端发送'C'字符表示准备就绪
            start_time = time.time()
            while (time.time() - start_time) < (self.timeout / 1000.0) and self.running:
                if self.ack_response == self.C:
                    self.waiting_for_ack = False
                    self.status_update.emit("接收端已准备就绪")
                    return True
                self.msleep(10)
                
            self.waiting_for_ack = False
            self.retry_count += 1
            self.status_update.emit(f"等待超时，重试 {self.retry_count}/{self.max_retries}")
            
        self.status_update.emit("接收端未在规定时间内准备就绪")
        return False
        
    def send_file_info(self):
        """发送文件信息包(包含文件名和文件大小)"""
        self.status_update.emit("发送文件信息包...")
        
        filename = os.path.basename(self.file_path).encode('ascii')
        filesize = str(self.file_size).encode('ascii')
        
        # 标准YModem格式: filename<NULL>filesize<SPACE>...<NULL padding>
        info_data = filename + b'\x00' + filesize + b' '
        
        # 用NULL填充剩余空间
        info_packet = info_data + b'\x00' * (self.packet_size - len(info_data))
        
        return self.send_packet(0, info_packet)
        
    def send_file_data(self):
        """发送文件数据包"""
        self.status_update.emit("开始发送文件数据...")
        
        for packet_num in range(1, self.total_packets + 1):
            if not self.running:
                return False
                
            # 计算数据包偏移和大小
            offset = (packet_num - 1) * self.packet_size
            packet_data = self.file_data[offset:offset + self.packet_size]
            
            # 如果数据不足一个包，用0x1A填充
            if len(packet_data) < self.packet_size:
                packet_data += b'\x1A' * (self.packet_size - len(packet_data))
                
            # 发送数据包
            if not self.send_packet(packet_num, packet_data):
                return False
                
            # 更新进度
            self.progress_update.emit(packet_num, self.total_packets)
            
        return True
        
    def send_packet(self, packet_num, data):
        """发送单个数据包"""
        self.retry_count = 0
        
        while self.retry_count < self.max_retries and self.running:
            # 构造数据包 - 根据配置的packet_size选择头部
            if self.packet_size == 128:
                header = self.SOH
            else:
                header = self.STX
                
            packet_num_byte = packet_num & 0xFF
            packet_num_inv = (~packet_num) & 0xFF
            
            # 计算CRC16校验
            crc = self.calculate_crc16(data)
            
            # 组装完整数据包
            packet = bytes([header, packet_num_byte, packet_num_inv]) + data + crc.to_bytes(2, 'big')
            
            # 发送数据包
            self.status_update.emit(f"发送数据包 {packet_num} (实际数据: {len(data)}字节, 总大小: {len(packet)}字节)")
            self.data_send.emit(packet)
            
            # 等待确认
            if packet_num == 0:
                # 文件信息包：等待ACK，C字符可选
                if self.wait_for_ack_or_ack_and_c():
                    self.status_update.emit(f"文件信息包发送成功")
                    return True
            else:
                # 数据包：只等待ACK
                if self.wait_for_ack():
                    self.status_update.emit(f"数据包 {packet_num} 发送成功")
                    return True
                
            self.retry_count += 1
            self.status_update.emit(f"数据包 {packet_num} 发送失败，重试 {self.retry_count}/{self.max_retries}")
            
        self.status_update.emit(f"数据包 {packet_num} 发送最终失败，已达最大重试次数")
        return False
        
    def wait_for_ack(self):
        """等待ACK确认"""
        self.waiting_for_ack = True
        self.ack_response = None
        
        self.status_update.emit("等待接收端ACK确认...")
        
        start_time = time.time()
        while (time.time() - start_time) < (self.timeout / 1000.0) and self.running:
            # 检查是否收到ACK响应
            if self.ack_response is not None:
                self.waiting_for_ack = False
                if self.ack_response == self.ACK:
                    self.status_update.emit("收到ACK确认")
                    return True
                elif self.ack_response == self.NAK:
                    self.status_update.emit("收到NAK，需要重传")
                    return False
                else:
                    self.status_update.emit(f"收到未知响应: 0x{self.ack_response:02X}")
                    return False
            self.msleep(10)
            
        self.waiting_for_ack = False
        self.status_update.emit("等待ACK超时")
        return False
        
    def wait_for_ack_or_ack_and_c(self):
        """等待ACK确认，C字符可选（用于文件信息包）"""
        self.status_update.emit("等待接收端ACK确认（C字符可选）...")
        
        self.waiting_for_ack = True
        self.ack_response = None
        ack_received = False
        c_received = False
        
        start_time = time.time()
        while (time.time() - start_time) < (self.timeout / 1000.0) and self.running:
            if self.ack_response == self.ACK and not ack_received:
                ack_received = True
                self.status_update.emit("收到ACK确认")
                # 等待一小段时间看是否还有C字符
                c_wait_start = time.time()
                self.ack_response = None  # 重置，等待可能的C字符
                
                # 等待额外的100ms看是否有C字符
                while (time.time() - c_wait_start) < 0.1 and self.running:
                    if self.ack_response == self.C:
                        c_received = True
                        self.status_update.emit("收到C字符")
                        break
                    self.msleep(10)
                
                # 无论是否收到C，都认为成功
                break
            elif self.ack_response == self.C and not ack_received:
                # 如果先收到C，检查是否有ACK在数据中
                if self.ACK in self.received_data[-5:]:
                    ack_received = True
                    c_received = True
                    self.status_update.emit("收到ACK + C确认")
                    break
            self.msleep(10)
            
        self.waiting_for_ack = False
        
        if ack_received:
            if c_received:
                self.status_update.emit("文件信息包确认完成：ACK + C")
            else:
                self.status_update.emit("文件信息包确认完成：仅ACK（兼容模式）")
            return True
        else:
            self.status_update.emit("未收到ACK确认")
            return False
        
    def wait_for_ack_and_c(self):
        """等待ACK + C确认（用于文件信息包）"""
        self.status_update.emit("等待接收端ACK + C确认...")
        
        self.waiting_for_ack = True
        self.ack_response = None
        ack_received = False
        c_received = False
        
        start_time = time.time()
        while (time.time() - start_time) < (self.timeout / 1000.0) and self.running:
            if self.ack_response == self.ACK and not ack_received:
                ack_received = True
                self.status_update.emit("收到ACK确认")
                self.ack_response = None  # 重置，等待下一个字符
            elif self.ack_response == self.C and ack_received and not c_received:
                c_received = True
                self.status_update.emit("收到C字符，可以继续发送数据")
                break
            elif self.ack_response == self.C and not ack_received:
                # 如果先收到C，可能是ACK+C一起快速到达，检查最近的数据
                if self.ACK in self.received_data[-5:]:
                    ack_received = True
                    c_received = True
                    self.status_update.emit("收到ACK + C确认")
                    break
            self.msleep(10)
            
        self.waiting_for_ack = False
        
        if ack_received and c_received:
            return True
        else:
            if not ack_received:
                self.status_update.emit("未收到ACK确认")
            elif not c_received:
                self.status_update.emit("未收到C字符")
            return False
        
    def send_end_of_transmission(self):
        """发送传输结束信号"""
        self.status_update.emit("发送传输结束信号...")
        
        # 第一次发送EOT
        self.status_update.emit("第一次发送EOT...")
        self.data_send.emit(bytes([self.EOT]))
        
        # 等待NAK（二次确认）
        if not self.wait_for_nak():
            self.status_update.emit("未收到NAK，传输可能有问题")
            return False
            
        # 第二次发送EOT
        self.status_update.emit("第二次发送EOT...")
        self.data_send.emit(bytes([self.EOT]))
        
        # 等待ACK,不是等待(ACK + C),提高兼容性
        if not self.wait_for_ack_or_ack_and_c():
            self.status_update.emit("未收到ACK or ACK+C，传输可能有问题")
            return False
        
        # 发送结束符包：SOH 00 FF NULL[128] CRC CRC
        self.status_update.emit("发送结束符包...")
        null_packet = b'\x00' * self.packet_size
        return self.send_final_packet(null_packet)
        
    def wait_for_nak(self):
        """等待NAK确认"""
        self.waiting_for_ack = True
        self.ack_response = None
        
        self.status_update.emit("等待接收端NAK确认...")
        
        start_time = time.time()
        while (time.time() - start_time) < (self.timeout / 1000.0) and self.running:
            if self.ack_response == self.NAK:
                self.waiting_for_ack = False
                self.status_update.emit("收到NAK确认")
                return True
            elif self.ack_response is not None:
                self.waiting_for_ack = False
                self.status_update.emit(f"收到非NAK响应: 0x{self.ack_response:02X}")
                return False
            self.msleep(10)
            
        self.waiting_for_ack = False
        self.status_update.emit("等待NAK超时")
        return False
        
    def send_final_packet(self, data):
        """发送最终结束符包"""
        header = self.SOH  # 结束符包使用SOH
        packet_num_byte = 0x00
        packet_num_inv = 0xFF
        
        # 计算CRC16校验
        crc = self.calculate_crc16(data)
        
        # 组装完整数据包
        packet = bytes([header, packet_num_byte, packet_num_inv]) + data + crc.to_bytes(2, 'big')
        
        # 发送数据包
        self.status_update.emit(f"发送结束符包 (总大小: {len(packet)}字节)")
        self.data_send.emit(packet)
        
        # 等待最终ACK确认
        if self.wait_for_ack():
            self.status_update.emit("结束符包发送成功")
            return True
        else:
            self.status_update.emit("结束符包发送失败")
            return False
        
    def calculate_crc16(self, data):
        """计算CRC16校验值"""
        crc = 0
        for byte in data:
            crc ^= byte << 8
            for _ in range(8):
                if crc & 0x8000:
                    crc = (crc << 1) ^ 0x1021
                else:
                    crc <<= 1
                crc &= 0xFFFF
        return crc
        
    def handle_received_data(self, data):
        """处理接收到的数据"""
        self.received_data += data
        
        # 如果正在等待ACK，检查最新接收的数据
        if self.waiting_for_ack and data:
            for byte in data:
                if byte in [self.ACK, self.NAK, self.CAN, self.C]:
                    # 发送ACK接收信号
                    self.ack_received.emit(byte)
                    self.status_update.emit(f"接收到响应字节: 0x{byte:02X}")
                    # 不要break，继续处理后续字节，支持ACK+C组合
        
        # 保留最近的几个字节用于判断响应
        if len(self.received_data) > 10:
            self.received_data = self.received_data[-10:]
            
    def on_ack_received(self, response_byte):
        """处理ACK接收信号"""
        if self.waiting_for_ack:
            self.ack_response = response_byte
            
    def stop(self):
        """停止协议传输"""
        self.running = False
        self.waiting_for_ack = False
        

class YmodemConfigWidget(QWidget):
    """Ymodem协议配置界面"""
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        
    def init_ui(self):
        """初始化配置界面"""
        layout = QVBoxLayout(self)
        
        # Ymodem协议参数组
        params_group = QGroupBox("Ymodem协议参数")
        params_layout = QGridLayout(params_group)
        
        # 数据包大小
        params_layout.addWidget(QLabel("数据包大小:"), 0, 0)
        self.packet_size_combo = QComboBox()
        self.packet_size_combo.addItems(["128", "1024"])
        self.packet_size_combo.setCurrentText("128")
        params_layout.addWidget(self.packet_size_combo, 0, 1)
        
        # 超时时间
        params_layout.addWidget(QLabel("超时时间(ms):"), 1, 0)
        self.timeout_spin = QSpinBox()
        self.timeout_spin.setRange(100, 10000)
        self.timeout_spin.setValue(1000)
        self.timeout_spin.setSuffix(" ms")
        params_layout.addWidget(self.timeout_spin, 1, 1)
        
        # 最大重试次数
        params_layout.addWidget(QLabel("最大重试次数:"), 2, 0)
        self.retry_spin = QSpinBox()
        self.retry_spin.setRange(1, 10)
        self.retry_spin.setValue(3)
        params_layout.addWidget(self.retry_spin, 2, 1)
        
        # CRC校验
        self.crc_check = QCheckBox("启用CRC16校验")
        self.crc_check.setChecked(True)
        params_layout.addWidget(self.crc_check, 3, 0, 1, 2)
        
        # 批量传输模式
        self.batch_check = QCheckBox("批量传输模式")
        self.batch_check.setChecked(False)
        params_layout.addWidget(self.batch_check, 4, 0, 1, 2)
        
        layout.addWidget(params_group)
        
    def get_config(self):
        """获取当前配置"""
        config = {
            'packet_size': self.packet_size_combo.currentText(),
            'timeout': str(self.timeout_spin.value()),
            'max_retries': str(self.retry_spin.value()),
            'crc_enabled': self.crc_check.isChecked(),
            'batch_mode': self.batch_check.isChecked()
        }
        # 调试输出配置信息
        print(f"DEBUG: GUI配置获取 - 数据包大小: {config['packet_size']}, 超时: {config['timeout']}, 重试: {config['max_retries']}")
        return config 